/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.lvyh.lightframe.job.admin.manager;

import com.google.common.collect.ImmutableSet;
import com.lvyh.lightframe.job.admin.schedule.ScheduleManager;
import com.lvyh.lightframe.job.admin.util.NetUtils;
import com.lvyh.lightframe.job.enums.ServerNodeChange;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.state.ConnectionState;
import org.apache.curator.framework.state.ConnectionStateListener;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * Schedule node cluster service
 */
@Service
public class ScheduleServerManager extends AbstractScheduleServerManager implements ConnectionStateListener {
    private final static Logger logger = LoggerFactory.getLogger(TaskNodeManager.class);
    /**
     * Server IP list
     */
    private volatile Set<String> serverList = new HashSet<String>();

    /**
     * Offline server IP list
     */
    private volatile Set<String> offlineList = new HashSet<String>();

    /**
     * Has the registered node been deactivated
     */
    private volatile boolean isRegisterOffline = false;

    @Autowired
    private ScheduleManager scheduleManager;

    private Watcher serverNodeWatcher = new Watcher() {
        @Override
        public void process(WatchedEvent event) {
            if (Event.EventType.NodeChildrenChanged.equals(event.getType())) {
                try {
                    logger.info("Reload service node, {}", event);
                    refreshServers();
                } finally {
                    watchServerNodes();
                }
            }
        }
    };

    private Watcher tmpJobNodeWatcher = new Watcher() {
        @Override
        public void process(WatchedEvent event) {
            if (Event.EventType.NodeChildrenChanged.equals(event.getType())) {
                try {
                    handleTmpJobs();
                } finally {
                    watchTmpJobNodes();
                }
            }
        }
    };

    private Watcher offlineNodeWatcher = new Watcher() {
        @Override
        public void process(WatchedEvent event) {
            if (Event.EventType.NodeChildrenChanged.equals(event.getType())) {
                try {
                    refreshOfflineServers();
                } finally {
                    watchOfflineNodes();
                }
            }
        }
    };

    @PostConstruct
    public void init() {

        logger.info("ServerNodeManager initialization...");
        try {
            zkManager.createPersistNode(SERVER_PATH, null);
            zkManager.createPersistNode(SERVER_TMP_JOBS, null);
            zkManager.createPersistNode(OFFLINE_SERVER_PATH, null);

            registerServerNode();

            registerOfflineServerNode();

            List<String> children = zkManager.getChildren(SERVER_PATH);
            logger.info("The current scheduling node is in offline state, loading the service node data:{}", SERVER_PATH);
            logger.info("Get online service node data:{}", children);
            serverList = new HashSet<>(children);

            watchServerNodes();

            watchOfflineNodes();

            this.triggerScheduleServerChangeListener(ServerNodeChange.ON_OFFLINE);

            // Add connection status event listener
            zkManager.addConnectionStateListener(this);

            this.triggerScheduleServerChangeListener(ServerNodeChange.ON_INIT);
            logger.info("ServerNodeManager Initialization complete!");

        } catch (Exception e) {
            e.printStackTrace();
        }

    }

    @PreDestroy
    public synchronized void destroy() {
        unRegisterServerNode();

        this.triggerScheduleServerChangeListener(ServerNodeChange.ON_DESTROY);
    }


    /**
     * Deactivate node  1.unRegisterServerNode 2.registerOffline
     */
    public synchronized void offlineServer() {

        this.unRegisterServerNode();

        this.registerOfflineServerNode();

        //Trigger disable listener
        this.triggerScheduleServerChangeListener(ServerNodeChange.ON_OFFLINE);

    }

    /**
     * Enable node
     */
    @Override
    public synchronized void onlineServer() {

        logger.info("ServerNodeManager Enable node...");

        // Register itself in the node list of ZK
        // Pull the server list in full once
        refreshServers();

        watchOfflineNodes();

        this.unRegisterOffServerNode();

        //Trigger enable listener
        this.triggerScheduleServerChangeListener(ServerNodeChange.ON_ONLINE);
    }

    /**
     * Add temporary task node
     */
    public void addTmpTaskNode(int taskId, String executeNode) throws Exception {
        String taskNode = zkManager.buildPath(SERVER_TMP_JOBS, getTmpJobName(taskId, executeNode));
        zkManager.createPersistNode(taskNode, null);
    }

    private String getTmpJobName(int taskId, String executeNode) {
        return executeNode + "-" + String.valueOf(taskId);
    }


    private void watchServerNodes() {
        logger.info("Add service node monitoring:{}", SERVER_PATH);
        try {
            zkManager.watchChildren(SERVER_PATH, serverNodeWatcher);
        } catch (Exception e) {
            logger.error(e.getMessage(), e);
        }
    }

    private void watchTmpJobNodes() {
        logger.info("Add temporary task node monitoring:{}", SERVER_TMP_JOBS);
        try {
            zkManager.watchChildren(SERVER_TMP_JOBS, tmpJobNodeWatcher);
        } catch (Exception e) {
            logger.error(e.getMessage(), e);
        }
    }

    private void watchOfflineNodes() {
        logger.info("Add or deactivate task node monitoring:{}", OFFLINE_SERVER_PATH);
        try {
            zkManager.watchChildren(OFFLINE_SERVER_PATH, offlineNodeWatcher);
        } catch (Exception e) {
            logger.error(e.getMessage(), e);
        }
    }

    /**
     * Handle temporary tasks
     */
    private synchronized void handleTmpJobs() {
        try {
            String localIp = NetUtils.getLocalAddress().getHostAddress();
            List<String> tmpJobs = zkManager.getChildren(SERVER_TMP_JOBS);

            // Just take the first one, because if you add it to the scheduler and delete the node at the same time,
            // the nodechildrenchanged event of ZK will be triggered. Just handle it in a loop
            if (!tmpJobs.isEmpty()) {
                if (scheduleManager == null) {
                    return;
                }

                for (String tmpJob : tmpJobs) {
                    if (tmpJob.startsWith(localIp + "-")) {
                        try {
                            String taskId = tmpJob.split("-")[1];
                            scheduleManager.addSchedule(Integer.valueOf(taskId), true);
                        } finally {
                            String nodePath = zkManager.buildPath(SERVER_TMP_JOBS, tmpJob);
                            // Delete ZK node immediately after adding temporary task
                            zkManager.deleteNode(nodePath);
                        }
                    }
                }
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * Refresh node
     */
    public synchronized void refreshServers() {
        try {
            registerServerNode();

            handleTmpJobs();

            watchServerNodes();

            watchTmpJobNodes();

            logger.info("Load service node data:{}", SERVER_PATH);
            List<String> children = zkManager.getChildren(SERVER_PATH);
            Set<String> servers = new HashSet<String>(children);

            logger.info("Load offline service node data:{}", OFFLINE_SERVER_PATH);
            List<String> offline = zkManager.getChildren(OFFLINE_SERVER_PATH);
            this.offlineList = new HashSet<String>(offline);

            boolean changed = !CollectionUtils.isEqualCollection(servers, serverList);

            serverList = servers;
            logger.info("Loading service node data completed, servers:" + servers);

            if (changed) {
                triggerScheduleServerChangeListener(ServerNodeChange.DATA_CHANGED);
            }

        } catch (Exception e) {
            logger.error(e.getMessage(), e);
        }
    }

    private synchronized void refreshOfflineServers() {
        try {
            this.watchOfflineNodes();
            logger.info("Load offline service node data:{}", OFFLINE_SERVER_PATH);
            List<String> offline = zkManager.getChildren(OFFLINE_SERVER_PATH);
            this.offlineList = new HashSet<String>(offline);
        } catch (Exception e) {
            logger.error(e.getMessage(), e);
        }

    }

    /**
     * Create a server node and set up server data
     */
    public boolean registerServerNode() {

        String localIp = NetUtils.getHostAddress();
        if (StringUtils.isBlank(localIp)) {
            throw new RuntimeException("serverIp can't be blank");
        }
        String serverNode = zkManager.buildPath(SERVER_PATH, localIp);
        try {
            boolean success = zkManager.createEphemeralNode(serverNode, createServerNodeData());
            if (success) {
                logger.info("Successfully registered service node:{}", localIp);
            }

            zkManager.createPersistNode(SERVER_TMP_JOBS, null);
            zkManager.createPersistNode(OFFLINE_SERVER_PATH, null);

            return success;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }

    }

    public boolean unRegisterServerNode() {

        String localIp = NetUtils.getLocalAddress().getHostAddress();

        if (StringUtils.isBlank(localIp)) {
            throw new RuntimeException("serverIp can't be blank");
        }
        String serverNode = zkManager.buildPath(SERVER_PATH, localIp);
        try {
            boolean success = zkManager.deleteNode(serverNode);
            if (success) {
                logger.info("Unregister service node succeeded:{}", localIp);
            }
            return success;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * Create and deactivate server nodes and set up server data
     */
    public boolean registerOfflineServerNode() {

        if (isRegisterOffline) {
            logger.warn("offline server node already register");
            return true;
        }

        String localIp = NetUtils.getHostAddress();

        if (StringUtils.isBlank(localIp)) {
            throw new RuntimeException("register offline server node error,serverIp can't be blank");
        }
        String serverNode = zkManager.buildPath(OFFLINE_SERVER_PATH, localIp);
        try {
            boolean success = zkManager.createEphemeralNode(serverNode, createServerNodeData());
            if (success) {
                logger.info("Successfully registered offline service node:{}", localIp);
            }
            this.isRegisterOffline = true;
            return success;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }

    }

    public boolean unRegisterOffServerNode() {
        if (!isRegisterOffline) {
            logger.warn("offline server node already unregister");
            return true;
        }
        String localIp = NetUtils.getHostAddress();

        if (StringUtils.isBlank(localIp)) {
            throw new RuntimeException("unregister offline server node error,serverIp can't be blank");
        }
        String serverNode = zkManager.buildPath(OFFLINE_SERVER_PATH, localIp);
        try {
            boolean success = zkManager.deleteNode(serverNode);
            if (success) {
                logger.info("Successfully unregistered offline service node:{}", localIp);
            }
            this.isRegisterOffline = false;
            return success;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public Set<String> getCurrentServerNodes() {
        return ImmutableSet.copyOf(serverList);
    }

    public Set<String> getCurrentOfflineServerNodes() {
        return ImmutableSet.copyOf(this.offlineList);
    }

    @Override
    public void stateChanged(CuratorFramework client, ConnectionState newState) {
        logger.warn("Zookeeeper connection status is abnormal:" + newState);
        if (ConnectionState.LOST.equals(newState) || ConnectionState.READ_ONLY.equals(newState) || ConnectionState.SUSPENDED.equals(newState)) {
            // When the connection is lost (or read-only), the service node will be cleared,
            // which will cause all scheduling tasks of the machine to be suspended
            serverList = new HashSet<String>();
            triggerScheduleServerChangeListener(ServerNodeChange.DATA_CHANGED);

        } else if (ConnectionState.CONNECTED.equals(newState) || ConnectionState.RECONNECTED.equals(newState)) {
            if (this.isRegisterOffline) {
                this.registerOfflineServerNode();
            } else {
                this.refreshServers();
            }
        }
    }
}
